import torch
import torch.nn as nn
import torch.nn.functional as F


class RecModel(nn.Module):
    def __init__(
            self,
            n_user,  # 用户数
            n_news,  # 电影数
            n_age,  # 年龄段数
            n_job,  # 工作数
            n_news_type,  # 电影类型数目
            n_news_title_dict,  # 电影名字典大小
            n_news_title_maxlen,  # 最长的电影名长度
            cnn_kernel_size_list,  # 文本卷积核窗口宽度列表（文本卷积核长度默认等于embeding_size）
            cnn_kernel_num,  # 文本卷积核个数
            device,  # 设备
            embed_dim=32,
            fc_size=200
    ):
        super().__init__()

        self.fc_size = fc_size
        self.cnn_kernel_num = cnn_kernel_num
        self.embed_dim = embed_dim
        self.device = device

        # ----------------------------------------- 用户通道 ---------------------------------------
        # user embedding
        self.embedding_uid = nn.Embedding(num_embeddings=n_user, embedding_dim=embed_dim)
        self.embedding_gender = nn.Embedding(num_embeddings=2, embedding_dim=embed_dim // 2)
        self.embedding_age = nn.Embedding(num_embeddings=n_age, embedding_dim=embed_dim // 2)
        self.embedding_job = nn.Embedding(num_embeddings=n_job, embedding_dim=embed_dim // 2)

        # user embed to fc: the first fully-connected layers
        self.fc_uid = nn.Linear(in_features=embed_dim, out_features=embed_dim)
        self.fc_gender = nn.Linear(in_features=embed_dim // 2, out_features=embed_dim)
        self.fc_age = nn.Linear(in_features=embed_dim // 2, out_features=embed_dim)
        self.fc_job = nn.Linear(in_features=embed_dim // 2, out_features=embed_dim)

        # concated embeddings to fc: the second fully-connected layer (n_batch,128) => (n_batch,200)
        self.fc_user_combined = nn.Linear(in_features=4 * embed_dim, out_features=fc_size)

        # ------------------------------------------ 电影通道 --------------------------------------
        # news embeddings
        self.embedding_mid = nn.Embedding(num_embeddings=n_news, embedding_dim=embed_dim)
        self.embedding_mtype = nn.EmbeddingBag(
            num_embeddings=n_news_type,
            embedding_dim=embed_dim,
            padding_idx=0
        )

        self.fc_mid = nn.Linear(in_features=embed_dim, out_features=embed_dim)
        self.fc_mtype = nn.Linear(in_features=embed_dim, out_features=embed_dim)

        # (n_batch, 64) => (n_batch,200)
        self.fc_mid_mtype_combined = nn.Linear(in_features=embed_dim * 2, out_features=fc_size)

        # text CNN: 文本卷积网络
        self.embedding_mtitle = nn.Embedding(num_embeddings=n_news_title_dict, embedding_dim=embed_dim, padding_idx=0)

        # 电影名向量化后维度: (n_batch, n_news_titile_maxlen, embed_dim)
        self.news_title_CNN = [
            nn.Sequential(
                nn.Conv2d(in_channels=1, out_channels=cnn_kernel_num, kernel_size=(k, embed_dim)),
                nn.ReLU(),
                nn.MaxPool2d(kernel_size=(n_news_title_maxlen - k + 1, 1), stride=(1, 1))
            ).to(device) for k in cnn_kernel_size_list
        ]

        # news channel concat
        self.fc_news_combine = nn.Linear(embed_dim * 2 + cnn_kernel_num * len(cnn_kernel_size_list), fc_size)

        self.output_fc = nn.Linear(fc_size * 2, 1)

    def forward(self, user_input, news_input):
        # unzip train_data
        uid = user_input['uid']
        gender = user_input['gender']
        age = user_input['age']
        job = user_input['job']

        mid = news_input['mid']
        mtitle = news_input['mtitle']
        mtype = news_input['mtype']

        uid, gender, age, job, mid, mtitle, mtype = \
            uid.to(self.device), gender.to(self.device), age.to(self.device), job.to(self.device), mid.to \
                (self.device), mtitle.to(self.device), mtype.to(self.device)

        # user channel forward
        feature_uid = F.relu(self.fc_uid(self.embedding_uid(uid)))
        feature_gender = F.relu(self.fc_gender(self.embedding_gender(gender)))
        feature_age = F.relu(self.fc_age(self.embedding_age(age)))
        feature_job = F.relu(self.fc_job(self.embedding_job(job)))

        # feature user (n_batch, 1, 200)
        feature_user = torch.tanh(
            self.fc_user_combined(
                torch.cat([feature_uid, feature_gender, feature_age, feature_job], 2)
            )
        ).view(-1, 1, 200)

        # news channel forward
        feature_mid = F.relu(self.fc_mid(self.embedding_mid(mid)))
        feature_mtype = self.embedding_mtype(mtype)

        # text CNN forward 文本卷积网络正向传播
        feature_img = self.embedding_mtitle(mtitle)
        flatten_tensors = []
        for conv in self.news_title_CNN:
            flatten_tensors.append(conv(feature_img).view(-1, 1, self.cnn_kernel_num))

        feature_flattern_dropout = F.dropout(torch.cat(flatten_tensors, 2), p=0.5)  # n_batch x embed_dim

        # feature_news n_batch x 1 x 200
        feature_news = torch.tanh(
            self.fc_news_combine(
                torch.cat([feature_mid.view(-1, 1, self.embed_dim), feature_mtype.view(-1, 1, self.embed_dim),
                           feature_flattern_dropout], 2)
            )
        )

        # feature user/news combined
        # output = self.output_fc(torch.cat([feature_user, feature_news], 2)).view(-1, 1)
        output = feature_user @ feature_news.T
        return output, feature_user, feature_news
